#!/usr/bin/python
# -*-coding:utf-8-*-
import pandas as pd

### 底层读取数据的依赖（不提供）
from zbc_factor_lib.base.factors_library_base import NewRQFactorLib as DataReader

# TODO - 第个月的第几个星期几交易日，作为调仓日
def get_monthly_nth_weekday_trade_date_data(start_date=None,
                                            end_date=None,
                                            weekday = 1,
                                            n_th = 1,
                                            selection_date_to_rebalance_date_interval = 1,
                                            filter_date='rebalance_start'):
    '''
    每个月的第n个星期m交易日
    '''
    data_reader = DataReader()

    # TODO - 读取交易日数据
    processed_trade_date_data = data_reader.read_basic_data_table('processed_trade_date_data')

    processed_trade_date_data = pd.DataFrame({'rebalance_start':processed_trade_date_data['trade_date'].unique()})
    processed_trade_date_data.index = processed_trade_date_data['rebalance_start'].copy()
    processed_trade_date_data['weekday'] = processed_trade_date_data.index.weekday + 1
    processed_trade_date_data['yearmonth'] = processed_trade_date_data.index.strftime('%Y%m')

    processed_trade_date_data['selection_date'] = \
        processed_trade_date_data['rebalance_start'].shift(selection_date_to_rebalance_date_interval)

    monthly_selected_trade_date_list = []
    for ym in processed_trade_date_data['yearmonth'].unique():
        monthly_trade_date_data = processed_trade_date_data[processed_trade_date_data['yearmonth'] == ym].copy()

        monthly_trade_date_data = monthly_trade_date_data[monthly_trade_date_data['weekday'] == weekday]

        if monthly_trade_date_data.shape[0] < n_th:
            if monthly_trade_date_data.shape[0] == 0:
                continue
            else:
                monthly_trade_date_data = monthly_trade_date_data.iloc[[-1], :]
        else:
           monthly_trade_date_data = monthly_trade_date_data.iloc[[n_th-1], :]

        monthly_selected_trade_date_list.append(monthly_trade_date_data)

    monthly_selected_trade_date = pd.concat(monthly_selected_trade_date_list, axis=0)

    monthly_selected_trade_date['rebalance_end'] = monthly_selected_trade_date['selection_date'].shift(-1)
    monthly_selected_trade_date['next_rebalance_start'] = monthly_selected_trade_date['rebalance_start'].shift(-1)

    output_columns = ['selection_date',
                      'rebalance_start',
                      'rebalance_end',
                      'next_rebalance_start']

    monthly_selected_trade_date = monthly_selected_trade_date[output_columns]
    # monthly_selected_trade_date = monthly_selected_trade_date[
    #     ['selection_date',
    #      'rebalance_start',
    #      'rebalance_end',
    #      'next_rebalance_start']
    #      # 'weekday',
    #      # 'yearmonth']
    # ]

    # TODO - 转换格式
    for c in output_columns:
        monthly_selected_trade_date[c] = pd.to_datetime(monthly_selected_trade_date[c])

    if start_date is not None and end_date is not None:
        monthly_selected_trade_date = monthly_selected_trade_date[(monthly_selected_trade_date[filter_date] >= start_date) &
                                                                  (monthly_selected_trade_date[filter_date] <= end_date)]
    elif start_date is None and end_date is not None:
        monthly_selected_trade_date = monthly_selected_trade_date[monthly_selected_trade_date[filter_date] <= end_date]
    elif start_date is not None and end_date is None:
        monthly_selected_trade_date = monthly_selected_trade_date[monthly_selected_trade_date[filter_date] >= start_date]

    return monthly_selected_trade_date.reset_index(drop=True)

# TODO - 第一个星期一交易日（1月，5月和10月以第一个交易日为主），作为调仓日
def get_monthly_first_monday_trade_date(start_date=None,
                                        end_date=None,
                                        special_treat_holiday=True,
                                        filter_date='rebalance_start'):
    '''
    每个月的第一个星期一交易日
    1月，5月和10月以第一个交易日为主（主要是因为放假）
    '''

    # TODO - 读取交易日数据
    data_reader = DataReader()

    processed_trade_date_data = data_reader.read_basic_data_table('processed_trade_date_data')
    processed_trade_date_data = processed_trade_date_data.set_index('natural_date')
    processed_trade_date_data = processed_trade_date_data.drop_duplicates('trade_date', keep='last')
    processed_trade_date_data.index = processed_trade_date_data['trade_date'].copy()
    processed_trade_date_data['weekday'] = pd.DatetimeIndex(processed_trade_date_data['trade_date']).weekday + 1

    processed_trade_date_data = processed_trade_date_data.dropna()

    # def get_first_monday_trade_date(monthly_trade_data):
    #     monthly_trade_data = monthly_trade_data.copy()
    #
    #     monthly_trade_data = monthly_trade_data[monthly_trade_data['weekday'] == 1]
    #     monthly_trade_data = monthly_trade_data[['trade_date']]
    #
    #     return monthly_trade_data.iloc[0]

    # 每月月底
    # monthly_resample = processed_trade_date_data.resample('1M')

    processed_trade_date_data['ym'] = processed_trade_date_data.index.strftime('%Y%m')

    monthly_first_monday_trade_date = processed_trade_date_data[processed_trade_date_data['weekday'] == 1]
    monthly_first_monday_trade_date = monthly_first_monday_trade_date.drop_duplicates('ym', keep='first')
    monthly_first_trade_date = processed_trade_date_data.drop_duplicates('ym', keep='first')

    # monthly_first_monday_trade_date = monthly_resample.apply(get_first_monday_trade_date)[['trade_date']]
    # monthly_first_trade_date = monthly_resample[['trade_date']].first()

    # TODO - 对1, 5, 10月进行特殊处理
    if special_treat_holiday:
        monthly_first_monday_trade_date['month'] = monthly_first_monday_trade_date.index.month
        monthly_first_trade_date['month'] = monthly_first_trade_date.index.month

        monthly_first_monday_trade_date = monthly_first_monday_trade_date[~monthly_first_monday_trade_date['month'].isin([1, 5, 10])]
        monthly_first_trade_date = monthly_first_trade_date[monthly_first_trade_date['month'].isin([1, 5, 10])]

        monthly_trade_date_data = pd.concat([monthly_first_monday_trade_date, monthly_first_trade_date],
                                            axis=0)

        monthly_trade_date_data = monthly_trade_date_data.sort_index()
    else:
        monthly_first_monday_trade_date['month'] = monthly_first_monday_trade_date.index.month

        monthly_trade_date_data = monthly_first_monday_trade_date.sort_index()

    # monthly_first_monday_trade_date = monthly_first_monday_trade_date.rename(columns={'trade_date':'first_monday_trade_date'})
    # monthly_first_trade_date = monthly_first_trade_date.rename(columns={'trade_date':'first_trade_date'})

    monthly_trade_date_data = monthly_trade_date_data.rename(columns={'trade_date':'rebalance_start',
                                                                      'last_trade_date': 'selection_date'})

    # monthly_trade_date_data = pd.concat([monthly_first_monday_trade_date, monthly_first_trade_date],
    #                                     axis=1)
    # monthly_trade_date_data['month'] = monthly_trade_date_data.index.month

    # monthly_trade_date_data['rebalance_start'] = monthly_trade_date_data['first_monday_trade_date'].copy()

    # start前一个交易日
    # monthly_trade_date_data['selection_date'] = \
    #     processed_trade_date_data.loc[monthly_trade_date_data['rebalance_start'], 'last_trade_date'].values

    monthly_trade_date_data['rebalance_end'] = monthly_trade_date_data['selection_date'].shift(-1)
    monthly_trade_date_data['next_rebalance_start'] = monthly_trade_date_data['rebalance_start'].shift(-1)
    # monthly_trade_date_data['end'] = monthly_resample['trade_date'].last()         # 月末最后一个交易日

    output_columns = ['selection_date',
                      'rebalance_start',
                      'rebalance_end',
                      'next_rebalance_start']

    monthly_trade_date_data = monthly_trade_date_data[output_columns]

    # TODO - 转换格式
    for c in output_columns:
        monthly_trade_date_data[c] = pd.to_datetime(monthly_trade_date_data[c])

    if start_date is not None and end_date is not None:
        monthly_trade_date_data = monthly_trade_date_data[(monthly_trade_date_data[filter_date] >= start_date) &
                                                          (monthly_trade_date_data[filter_date] <= end_date)]
    elif start_date is not None and end_date is None:
        monthly_trade_date_data = monthly_trade_date_data[monthly_trade_date_data[filter_date] >= start_date]
    elif start_date is None and end_date is not None:
        monthly_trade_date_data = monthly_trade_date_data[monthly_trade_date_data[filter_date] <= end_date]

    return monthly_trade_date_data.reset_index(drop=True)

def get_monthly_first_trade_date(start_date=None,
                                  end_date=None,
                                  filter_date='rebalance_start'):
    '''
    每个月的第一个交易日
    '''

    # TODO - 读取交易日数据
    data_reader = DataReader()

    processed_trade_date_data = data_reader.read_basic_data_table('processed_trade_date_data')
    processed_trade_date_data = processed_trade_date_data.set_index('natural_date')
    processed_trade_date_data = processed_trade_date_data.drop_duplicates('trade_date', keep='last')
    processed_trade_date_data.index = processed_trade_date_data['trade_date'].copy()
    processed_trade_date_data['weekday'] = pd.DatetimeIndex(processed_trade_date_data['trade_date']).weekday + 1

    processed_trade_date_data = processed_trade_date_data.dropna()

    processed_trade_date_data['ym'] = processed_trade_date_data.index.strftime('%Y%m')

    monthly_first_trade_date = processed_trade_date_data.drop_duplicates('ym', keep='first')

    monthly_trade_date_data = monthly_first_trade_date.sort_index()

    monthly_trade_date_data = monthly_trade_date_data.rename(columns={'trade_date':'rebalance_start',
                                                                      'last_trade_date': 'selection_date'})

    monthly_trade_date_data['rebalance_end'] = monthly_trade_date_data['selection_date'].shift(-1)
    monthly_trade_date_data['next_rebalance_start'] = monthly_trade_date_data['rebalance_start'].shift(-1)

    output_columns = ['selection_date',
                      'rebalance_start',
                      'rebalance_end',
                      'next_rebalance_start']

    monthly_trade_date_data = monthly_trade_date_data[output_columns]

    # TODO - 转换格式
    for c in output_columns:
        monthly_trade_date_data[c] = pd.to_datetime(monthly_trade_date_data[c])

    if start_date is not None and end_date is not None:
        monthly_trade_date_data = monthly_trade_date_data[(monthly_trade_date_data[filter_date] >= start_date) &
                                                          (monthly_trade_date_data[filter_date] <= end_date)]
    elif start_date is not None and end_date is None:
        monthly_trade_date_data = monthly_trade_date_data[monthly_trade_date_data[filter_date] >= start_date]
    elif start_date is None and end_date is not None:
        monthly_trade_date_data = monthly_trade_date_data[monthly_trade_date_data[filter_date] <= end_date]

    return monthly_trade_date_data.reset_index(drop=True)

# TODO - 每个星期几，作为调仓日
def get_every_weekday_trade_date(start_date=None,
                                 end_date=None,
                                 filter_date='rebalance_start',
                                 special_treat_holiday=True,
                                 minimum_holding_days=1.0,
                                 weekday=1):
    '''
    每个月的第一个星期一交易日
    1月，5月和10月以第一个交易日为主（主要是因为放假）
    '''
    data_reader = DataReader()

    # TODO - 读取交易日数据
    processed_trade_date_data = data_reader.read_basic_data_table('processed_trade_date_data')
    processed_trade_date_data = processed_trade_date_data.set_index('natural_date')
    processed_trade_date_data = processed_trade_date_data.drop_duplicates('trade_date', keep='last')
    processed_trade_date_data.index = processed_trade_date_data['trade_date'].copy()
    processed_trade_date_data['weekday'] = pd.DatetimeIndex(processed_trade_date_data['trade_date']).weekday + 1

    processed_trade_date_data = processed_trade_date_data.dropna()

    weekday_trade_date_data = processed_trade_date_data[processed_trade_date_data['weekday'] == weekday]

    weekday_trade_date_data['rebalance_end'] = weekday_trade_date_data['last_trade_date'].shift(-1)
    weekday_trade_date_data['next_rebalance_start'] = weekday_trade_date_data['trade_date'].shift(-1)

    weekday_trade_date_data = weekday_trade_date_data.rename(columns={'trade_date':'rebalance_start',
                                                                    'last_trade_date':'selection_date'})


    output_columns = ['selection_date',
                      'rebalance_start',
                      'rebalance_end',
                      'next_rebalance_start']

    weekday_trade_date_data = weekday_trade_date_data[output_columns]

    for c in output_columns:
        weekday_trade_date_data[c] = pd.to_datetime(weekday_trade_date_data[c])

    weekday_trade_date_data = weekday_trade_date_data.reset_index(drop=True)

    weekday_trade_date_data = weekday_trade_date_data.iloc[:-1, :]

    if start_date is not None and end_date is not None:
        filter_end_date = pd.to_datetime(end_date) + pd.Timedelta(days=30)
        filter_end_date = filter_end_date.strftime('%Y-%m-%d')
        weekday_trade_date_data = weekday_trade_date_data[(weekday_trade_date_data[filter_date] >= start_date) &
                                                          (weekday_trade_date_data[filter_date] <= filter_end_date)]
    elif start_date is not None and end_date is None:
        weekday_trade_date_data = weekday_trade_date_data[weekday_trade_date_data[filter_date] >= start_date]
    elif start_date is None and end_date is not None:
        filter_end_date = pd.to_datetime(end_date) + pd.Timedelta(days=30)
        filter_end_date = filter_end_date.strftime('%Y-%m-%d')
        weekday_trade_date_data = weekday_trade_date_data[weekday_trade_date_data[filter_date] <= filter_end_date]

    weekday_trade_date_data['weekday'] = pd.DatetimeIndex(weekday_trade_date_data['rebalance_start']).weekday + 1

    ### TODO - 专门处理1,5,10三个月的交易日数据
    if special_treat_holiday:
        # def get_first_monday_trade_date(monthly_trade_data):
        #     monthly_trade_data = monthly_trade_data.copy()
        #
        #     monthly_trade_data = monthly_trade_data[monthly_trade_data['weekday'] == 1]
        #     monthly_trade_data = monthly_trade_data[['trade_date', 'last_trade_date']]
        #
        #     # return [monthly_trade_data.iloc[0, 0], monthly_trade_data.iloc[0, 1]]
        #     return monthly_trade_data.iloc[0, :]

        # 每月月底
        processed_trade_date_data['ym'] = processed_trade_date_data.index.strftime('%Y%m')
        monthly_first_monday_trade_date = processed_trade_date_data[processed_trade_date_data['weekday'] == 1]

        monthly_first_monday_trade_date = monthly_first_monday_trade_date.drop_duplicates('ym', keep='first')

        monthly_resample = processed_trade_date_data.resample('1M')

        # monthly_first_monday_trade_date = \
        #     monthly_resample.apply(get_first_monday_trade_date)[['trade_date', 'last_trade_date']]

        monthly_first_trade_date = monthly_resample[['trade_date', 'last_trade_date']].first()

        monthly_first_monday_trade_date = \
            monthly_first_monday_trade_date.rename(columns={'trade_date':'first_monday_trade_date',
                                                            'last_trade_date': 'first_last_monday_trade_date'})

        monthly_first_trade_date = monthly_first_trade_date.rename(columns={'trade_date':'first_trade_date',
                                                                            'last_trade_date':'first_last_trade_date'})

        monthly_first_monday_trade_date.index = monthly_first_trade_date.index

        monthly_trade_date = pd.concat([monthly_first_monday_trade_date, monthly_first_trade_date], axis=1)
        monthly_trade_date['month'] = monthly_trade_date.index.month

        monthly_trade_date = monthly_trade_date[monthly_trade_date['month'].isin([1, 5, 10])]

        monthly_trade_date['weekday'] = pd.DatetimeIndex(monthly_trade_date['first_trade_date']).weekday + 1
        monthly_trade_date = monthly_trade_date.rename(columns={'first_trade_date': 'rebalance_start',
                                                                'first_last_trade_date': 'selection_date',
                                                                'first_last_monday_trade_date': 'rebalance_end',
                                                                'first_monday_trade_date':'next_rebalance_start'})

        monthly_trade_date = \
            monthly_trade_date[monthly_trade_date['selection_date'] != monthly_trade_date['rebalance_end']]

        monthly_trade_date = monthly_trade_date.reset_index(drop=True)
        monthly_trade_date = monthly_trade_date[weekday_trade_date_data.columns]

        if start_date is not None and end_date is not None:
            filter_end_date = pd.to_datetime(end_date) + pd.Timedelta(days=30)
            filter_end_date = filter_end_date.strftime('%Y-%m-%d')
            monthly_trade_date = monthly_trade_date[(monthly_trade_date[filter_date] >= start_date) &
                                                    (monthly_trade_date[filter_date] <= filter_end_date)]
        elif start_date is not None and end_date is None:
            monthly_trade_date = monthly_trade_date[monthly_trade_date[filter_date] >= start_date]
        elif start_date is None and end_date is not None:
            filter_end_date = pd.to_datetime(end_date) + pd.Timedelta(days=30)
            filter_end_date = filter_end_date.strftime('%Y-%m-%d')
            monthly_trade_date = monthly_trade_date[monthly_trade_date[filter_date] <= filter_end_date]

        weekday_trade_date_data = pd.concat([weekday_trade_date_data, monthly_trade_date], axis=0)
        weekday_trade_date_data = weekday_trade_date_data.sort_values('selection_date')

        weekday_trade_date_data['rebalance_end'] = weekday_trade_date_data['selection_date'].shift(-1)
        weekday_trade_date_data['next_rebalance_start'] = weekday_trade_date_data['rebalance_start'].shift(-1)

        # TODO - 过滤持仓少于minimum_holding_days+1天的情况
        weekday_trade_date_data = \
            weekday_trade_date_data[
                weekday_trade_date_data['rebalance_start']  + pd.Timedelta(days=minimum_holding_days) <= weekday_trade_date_data['rebalance_end']
            ]

        weekday_trade_date_data['rebalance_end'] = weekday_trade_date_data['selection_date'].shift(-1)
        weekday_trade_date_data['next_rebalance_start'] = weekday_trade_date_data['rebalance_start'].shift(-1)

        if start_date is not None and end_date is not None:
            weekday_trade_date_data = weekday_trade_date_data[(weekday_trade_date_data[filter_date] >= start_date) &
                                                              (weekday_trade_date_data[filter_date] <= end_date)]
        elif start_date is not None and end_date is None:
            weekday_trade_date_data = weekday_trade_date_data[weekday_trade_date_data[filter_date] >= start_date]
        elif start_date is None and end_date is not None:
            weekday_trade_date_data = weekday_trade_date_data[weekday_trade_date_data[filter_date] <= end_date]

    return weekday_trade_date_data.dropna().reset_index(drop=True)

# TODO - 每N个交易日，作为调仓日
def get_every_fixed_frequency_trade_date(start_date=None,
                                         end_date=None,
                                         freq=5,
                                         filter_date='rebalance_start'):
    '''
    以固定的间隔交易日来生成交易日期
    '''

    data_reader = DataReader()

    # TODO - 读取交易日数据
    processed_trade_date_data = data_reader.read_basic_data_table('processed_trade_date_data')
    processed_trade_date_data = processed_trade_date_data.set_index('natural_date')
    processed_trade_date_data = processed_trade_date_data.drop_duplicates('trade_date', keep='last')
    processed_trade_date_data.index = processed_trade_date_data['trade_date'].copy()

    fixed_interval_trade_date_data = processed_trade_date_data.copy()

    fixed_interval_trade_date_data['rebalance_start'] = fixed_interval_trade_date_data['trade_date'].copy()
    fixed_interval_trade_date_data['selection_date'] = fixed_interval_trade_date_data['last_trade_date'].copy()

    fixed_interval_trade_date_data = fixed_interval_trade_date_data.iloc[range(0, fixed_interval_trade_date_data.shape[0], freq)]

    fixed_interval_trade_date_data['rebalance_end'] = fixed_interval_trade_date_data['selection_date'].shift(-1)
    fixed_interval_trade_date_data['next_rebalance_start'] = fixed_interval_trade_date_data['rebalance_start'].shift(-1)

    output_columns = ['selection_date',
                      'rebalance_start',
                      'rebalance_end',
                      'next_rebalance_start']

    fixed_interval_trade_date_data = fixed_interval_trade_date_data[output_columns]

    # TODO - 转换格式
    for c in output_columns:
        fixed_interval_trade_date_data[c] = pd.to_datetime(fixed_interval_trade_date_data[c])

    if start_date is not None and end_date is not None:
        fixed_interval_trade_date_data = fixed_interval_trade_date_data[(fixed_interval_trade_date_data[filter_date] >= start_date) &
                                                          (fixed_interval_trade_date_data[filter_date] <= end_date)]
    elif start_date is not None and end_date is None:
        fixed_interval_trade_date_data = fixed_interval_trade_date_data[fixed_interval_trade_date_data[filter_date] >= start_date]
    elif start_date is None and end_date is not None:
        fixed_interval_trade_date_data = fixed_interval_trade_date_data[fixed_interval_trade_date_data[filter_date] <= end_date]

    return fixed_interval_trade_date_data.reset_index(drop=True)


if __name__ == '__main__':
    # start_date = '2017-01-01'
    # start_date = '2010-01-01'
    start_date = '2020-01-03'
    end_date = '2020-12-31'

    # selected_trade_date_data = get_monthly_nth_weekday_trade_date_data(start_date=start_date,
    #                                                                    end_date=end_date,
    #                                                                    weekday=1,
    #                                                                    n_th=2,
    #                                                                    selection_date_to_rebalance_date_interval=1,
    #                                                                    filter_date='rebalance_start')
    #
    # selected_trade_date_data = get_monthly_first_monday_trade_date(start_date=start_date,
    #                                                                end_date=end_date,
    #                                                                filter_date='rebalance_start')

    selected_trade_date_data = get_every_weekday_trade_date(start_date=start_date,
                                                            end_date=end_date,
                                                            weekday=1,
                                                            minimum_holding_days=3.0,
                                                            filter_date='rebalance_start')

    # selected_trade_date_data.iloc[::2, :]

    # selected_trade_date_data = get_every_fixed_frequency_trade_date(start_date=start_date,
    #                                                                 end_date=end_date,
    #                                                                 freq=2,
    #                                                                 filter_date='rebalance_start')




